import os
import time
import acl
import numpy as np
import av
import cv2
import constant as Const

from termcolor import cprint
from himpi_util import check_ret, align, get_video_total_frames

class Vdec(object):
	def __init__(self, device_id, file_path):
		self.device_id = device_id
		self.file_path = file_path
		self.context = None
		self.stream = None
		self.timeout = 1000
		self.send_count = 0
		self.get_count = 0
		self.init_resource()
		self.create()

	def __del__(self):
		print("[Vedc] release resource:")
	
	def init_resource(self):
		print("[Vdec] init resource stage:")

		ret = acl.init()
		check_ret("acl.init", ret)

		ret = acl.rt.set_device(self.device_id)
		check_ret("acl.rt.set_device", ret)

		self.context, ret = acl.rt.create_context(self.device_id)
		check_ret("acl.rt.create_context", ret)

		ret = acl.himpi.sys_init()
		check_ret("acl.himpi.sys_init", ret)

		self.init_file()

		width_stride = align(self.video_stream.width, 16) * 3
		height_stride = self.video_stream.height
		self.buffer_size = width_stride * height_stride
		self.out_ptr, ret = acl.rt.malloc_host(self.buffer_size)
		check_ret("acl.rt.malloc_host", ret)


		self.in_buffer, ret = acl.himpi.dvpp_malloc(self.device_id, self.buffer_size)
		check_ret("acl.himpi.dvpp_malloc", ret)

		self.out_buffer, ret = acl.himpi.dvpp_malloc(self.device_id, self.buffer_size)
		check_ret("acl.himpi.dvpp_malloc", ret)

		self.vdec_stream = {
			"end_of_frame": 1,
			"end_of_stream": 0,
			"need_display": 1,
			"pts": time.time_ns(),
			"private_data": 0,
			"len": 0,
			"addr": self.in_buffer}
		self.vdec_pic_info = {
			"width": self.video_stream.width,
			"height": self.video_stream.height,
			"width_stride": width_stride,
			"height_stride": height_stride,
			"pixel_format": Const.HI_PIXEL_FORMAT_RGB_888,
			"vir_addr": self.out_buffer,
			"buffer_size": self.buffer_size}

		width = align(self.video_stream.width, 16)
		height = align(self.video_stream.height, 2)

		tmv_buf_size = acl.himpi.vdec_get_tmv_buf_size(Const.HI_PT_H264, width, height)
		self.attr = {
			"type": Const.HI_PT_H264, 
			"mode": 1,
			"pic_width": width,
			"pic_height": height,
			"stream_buf_size": width * height * 4,
			"frame_buf_cnt": 0,
			"video_attr": {
				"ref_frame_num": 5,
				"temporal_mvp_en": 1,
				"tmv_buf_size": tmv_buf_size}}
		print("[Vdec] init resource stage success")		


	def init_file(self):
		self.total_frame = 50
		self.video = av.open(self.file_path, options={'rtsp_transport': 'tcp'})
		self.video_stream = [s for s in self.video.streams if s.type == "video"][0]


	def create(self):
		self.chn = 0
		ret = acl.himpi.vdec_create_chn(self.chn, self.attr)
		check_ret("acl.himpi.vdec_create_chn", ret)

		param, ret = acl.himpi.vdec_get_chn_param(self.chn)
		check_ret("acl.himpi.vdec_get_chn_param", ret)

		param["video_param"]["dec_mode"] = Const.HI_VIDEO_DEC_MODE_IPB
		param["video_param"]["compress_mode"] = Const.HI_COMPRESS_MODE_HFBC
		param["video_param"]["video_format"] = Const.HI_VIDEO_FORMAT_TILE_64x16
		param["display_frame_num"] = 2
		g_output_order = 0
		if (g_output_order == 0):
			param["video_param"]["out_order"] = Const.HI_VIDEO_OUT_ORDER_DISPLAY
		else:
			param["video_param"]["out_order"] = Const.HI_VIDEO_OUT_ORDER_DEC


		ret = acl.himpi.vdec_set_chn_param(self.chn, param)
		check_ret("acl.himpi.vdec_set_chn_param", ret)

		ret = acl.himpi.vdec_start_recv_stream(self.chn)
		check_ret("acl.himpi.vdec_start_recv_stream", ret)


	def vdec_send_frame(self):
		print(f"Chn {self.chn} send_stream Thread Start")
		ret = acl.rt.set_context(self.context)
		check_ret("acl.rt.set_context", ret)

		count = 0
		for packet in self.video.demux([self.video_stream]):
			size = packet.buffer_size
			ret = acl.rt.memcpy(self.vdec_stream["addr"], size, packet.buffer_ptr, size, Const.ACL_MEMCPY_HOST_TO_DEVICE)
			check_ret("acl.rt.memcpy", ret)

			self.vdec_stream["len"] = size
			ret = Const.HI_ERR_VDEC_BUF_FULL
			while (ret == Const.HI_ERR_VDEC_BUF_FULL):
				ret = acl.himpi.vdec_send_stream(self.chn, self.vdec_stream, self.vdec_pic_info, self.timeout)
				self.send_count += 1
				if (self.send_count - self.get_count > 15):
					time.sleep(0.1)
					print("---------------")
			check_ret("acl.himpi.vdec_send_stream", ret)
			count += 1
			if count == self.total_frame:
				break
		
		vdec_stream = {
			"end_of_frame": 0, 
			"end_of_stream": 1,
			"len": 0,
			"need_display": 1,
			"pts": time.time_ns(),
			"addr": 0}
		picture_info = {
			"vir_addr": 0,
			"buffer_size": 0}

		ret = acl.himpi.vdec_send_stream(self.chn, vdec_stream, picture_info, -1)
		check_ret("acl.himpi.vdec_send_stream", ret)
		print(f"Chn {self.chn} send_frame Thread Exit")
	
	def get_frame(self):
		print(f"Chn {self.chn} get_frame Thread Start")
		ret = acl.rt.set_context(self.context)
		check_ret("acl.rt.setcontext", ret)
		count = 0
		while(count < self.total_frame):
			frame_info, supplement, stram, ret = acl.himpi.vdec_get_frame(self.chn, self.timeout)
			self.get_count += 1
			if (ret != 0):
				time.sleep(0.1)
				print("no buffer:")
			else:
				count += 1
				print(f"Chn {self.chn} GetFrame Success, Decode Success[{count}]")
			ret = acl.himpi.vdec_release_frame(self.chn, frame_info)
		print(f"Chn {self.chn} get_frame Thred Exit")
	
